"""
基础数据模型
"""

from datetime import datetime
from typing import Optional, List, Dict, Any, Union
from pydantic import BaseModel, Field
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage, ToolMessage


class Message(BaseModel):
    """消息模型"""
    role: str = Field(..., description="消息角色，如'user'、'assistant'、'system'、'tool'等")
    content: str = Field(..., description="消息内容")
    tool_call_id: Optional[str] = Field(None, description="工具调用ID (仅当 role='tool' 时使用)")
    
    def to_langchain(self) -> BaseMessage:
        """将消息转换为LangChain消息格式"""
        if self.role == "user":
            return HumanMessage(content=self.content)
        elif self.role == "assistant":
            return AIMessage(content=self.content)
        elif self.role == "system":
            return SystemMessage(content=self.content)
        elif self.role == "tool":
            if not self.tool_call_id:
                raise ValueError("ToolMessage 必须有关联的 tool_call_id")
            return ToolMessage(content=self.content, tool_call_id=self.tool_call_id)
        else:
            return HumanMessage(content=f"[{self.role}] {self.content}")

    class Config:
        json_schema_extra = {
            "example": {
                "role": "user",
                "content": "将'夏日凉鞋'价格改为199元"
            }
        }


class ChatRequest(BaseModel):
    """聊天请求模型"""
    messages: List[Message]
    session_id: Optional[str] = None

    model_config = {
        "json_schema_extra": {
            "example": {
                "session_id": "sess_1688626800",
                "messages": [
                    {
                        "role": "user",
                        "content": "你好，我想了解一下如何提升我的商品曝光率。"
                    }
                ]
            }
        }
    }


class ToolCall(BaseModel):
    """工具调用模型"""
    name: str = Field(..., description="工具名称")
    arguments: Dict[str, Any] = Field(..., description="工具参数")
    
    class Config:
        json_schema_extra = {
            "example": {
                "name": "update_product_price",
                "arguments": {
                    "product_name": "夏日凉鞋",
                    "new_price": 199
                }
            }
        }


class AgentResponse(BaseModel):
    """Agent响应模型"""
    message: str = Field(..., description="响应消息")
    tool_calls: Optional[List[ToolCall]] = Field(None, description="工具调用")
    
    class Config:
        json_schema_extra = {
            "example": {
                "message": "我将帮您将'夏日凉鞋'的价格修改为199元",
                "tool_calls": [
                    {
                        "name": "update_product_price",
                        "arguments": {
                            "product_name": "夏日凉鞋",
                            "new_price": 199
                        }
                    }
                ]
            }
        }


class TaskStatus(BaseModel):
    """任务状态模型"""
    task_id: str = Field(..., description="任务ID")
    status: str = Field(..., description="任务状态: pending, running, completed, failed")
    created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
    updated_at: datetime = Field(default_factory=datetime.now, description="更新时间")
    result: Optional[Dict[str, Any]] = Field(None, description="任务结果")
    error: Optional[str] = Field(None, description="错误信息")
    
    class Config:
        json_schema_extra = {
            "example": {
                "task_id": "task_123456",
                "status": "completed",
                "created_at": "2023-11-01T12:00:00",
                "updated_at": "2023-11-01T12:01:30",
                "result": {
                    "message": "价格已成功更新为199元"
                },
                "error": None
            }
        } 


class ToolOutput(BaseModel):
    """单个工具调用的输出模型"""
    tool_call_id: str = Field(..., description="工具调用的唯一ID")
    output: str = Field(..., description="工具执行返回的字符串结果")


class ToolResumeRequest(BaseModel):
    """恢复工作流的请求模型"""
    session_id: str
    tool_outputs: List[ToolOutput]

    model_config = {
        "json_schema_extra": {
            "example": {
                "session_id": "sess_1688626800",
                "tool_outputs": [
                    {
                        "tool_call_id": "tool_abc123",
                        "output": "{\"success\": true, \"message\": \"成功导航到 https://example.com\"}"
                    }
                ]
            }
        }
    } 